import torch
import torchvision.models as models

# 保存和加载模型权重 

# PyTorch 模型将学习到的参数存储在一个名为 state_dict 的内部状态字典中。可以通过 torch.save 方法持久化这些参数。

model = models.vgg16(weights='IMAGENET1K_V1')
torch.save(model.state_dict(), '/home/misaka/AI-self-learning/weight/model_weights.pth')

# 要加载模型权重，需要先创建相同模型的实例，然后使用 load_state_dict() 方法加载参数。

model = models.vgg16() # 我们不指定“权重”，即创建未经训练的模型
model.load_state_dict(torch.load('/home/misaka/AI-self-learning/weight/model_weights.pth'))
model.eval()


# 保存和加载具有形状的模型
# 加载模型权重时，我们需要先实例化模型类，因为该类定义了网络的结构。
# 我们可能希望将此类的结构与模型一起保存，在这种情况下，我们可以将 model（而不是 model.state_dict()）传递给保存函数。

torch.save(model, '/home/misaka/AI-self-learning/weight/model.pth')
# 然后，我们可以使用 torch.load() 函数加载模型。
model = torch.load('/home/misaka/AI-self-learning/weight/model.pth')


